package com.ben.multiple.shard.springjdbc;

import com.google.common.base.Splitter;
import com.google.common.collect.Range;
import org.apache.shardingsphere.infra.util.exception.ShardingSpherePreconditions;
import org.apache.shardingsphere.infra.util.exception.external.sql.type.generic.UnsupportedSQLOperationException;
import org.apache.shardingsphere.sharding.algorithm.sharding.ShardingAutoTableAlgorithmUtils;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;
import org.apache.shardingsphere.sharding.exception.algorithm.sharding.ShardingAlgorithmInitializationException;
import org.apache.shardingsphere.sharding.exception.data.NullShardingValueException;

import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Properties;

/**
 * @author ben
 */
public class CustomIntegerStandardShardingAlgorithm implements StandardShardingAlgorithm<Integer> {
    private static final String SHARDING_RANGES_KEY = "sharding-ranges";

    private Map<Integer, Range<Integer>> partitionRange;

    @Override
    public void init(final Properties props) {
        partitionRange = calculatePartitionRange(props);
    }

    public Map<Integer, Range<Integer>> calculatePartitionRange(final Properties props) {
        ShardingSpherePreconditions.checkState(props.containsKey(SHARDING_RANGES_KEY), () -> new ShardingAlgorithmInitializationException(getType(), "Sharding ranges cannot be null."));
        List<Integer> partitionRanges = Splitter.on(",").trimResults().splitToList(props.getProperty(SHARDING_RANGES_KEY))
                .stream().map(this::parse).filter(Objects::nonNull).sorted().toList();
        ShardingSpherePreconditions.checkState(!partitionRanges.isEmpty(), () -> new ShardingAlgorithmInitializationException(getType(), "Sharding ranges can not be empty."));
        Map<Integer, Range<Integer>> result = new HashMap<>(partitionRanges.size() + 1, 1F);
        for (int i = 0; i < partitionRanges.size(); i++) {
            Integer rangeValue = partitionRanges.get(i);
            if (0 == i) {
                result.put(0, Range.lessThan(rangeValue));
            } else {
                Integer previousRangeValue = partitionRanges.get(i - 1);
                result.put(i, Range.closedOpen(previousRangeValue, rangeValue));
            }
            if (i == partitionRanges.size() - 1) {
                result.put(i + 1, Range.atLeast(rangeValue));
            }
        }
        return result;
    }

    private Integer parse(final String value) {
        try {
            return Integer.parseInt(value);
        } catch (final NumberFormatException ex) {
            return null;
        }
    }

    @Override
    public String getType() {
        return "CUSTOM-INT-STANDARD";
    }

    private Integer getPartition(final Comparable<?> value) {
        for (Map.Entry<Integer, Range<Integer>> entry : partitionRange.entrySet()) {
            if (entry.getValue().contains(getIntValue(value))) {
                return entry.getKey();
            }
        }
        throw new UnsupportedSQLOperationException("");
    }

    private Integer getIntValue(final Comparable<?> value) {
        return Integer.parseInt(value.toString());
    }

    private int getFirstPartition(final Range<Integer> valueRange) {
        return valueRange.hasLowerBound() ? getPartition(valueRange.lowerEndpoint()) : 0;
    }

    private int getLastPartition(final Range<Integer> valueRange) {
        return valueRange.hasUpperBound() ? getPartition(valueRange.upperEndpoint()) : partitionRange.size() - 1;
    }

    @Override
    public String doSharding(Collection<String> availableTargetNames, PreciseShardingValue<Integer> shardingValue) {
        ShardingSpherePreconditions.checkNotNull(shardingValue.getValue(), NullShardingValueException::new);
        ShardingSpherePreconditions.checkNotNull(shardingValue.getValue(), NullShardingValueException::new);
        String suffix = String.valueOf(getPartition(shardingValue.getValue()));
        return ShardingAutoTableAlgorithmUtils.findMatchedTargetName(availableTargetNames, suffix, shardingValue.getDataNodeInfo()).orElse(null);
    }

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, RangeShardingValue<Integer> shardingValue) {
        Collection<String> result = new LinkedHashSet<>(availableTargetNames.size(), 1F);
        int firstPartition = getFirstPartition(shardingValue.getValueRange());
        int lastPartition = getLastPartition(shardingValue.getValueRange());
        for (int partition = firstPartition; partition <= lastPartition; partition++) {
            String suffix = String.valueOf(partition);
            ShardingAutoTableAlgorithmUtils.findMatchedTargetName(availableTargetNames, suffix, shardingValue.getDataNodeInfo()).ifPresent(result::add);
        }
        return result;
    }
}
